#!/usr/bin/python3
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

import csv
import re
import os
import codecs

def webOfScience(filename,option='min',encoding='utf-8'):
  '''
# Load citation data from web of science format (tsv)
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  goto http://www.webofknowledge.com/
#  search your target
#  send to file (Save to Other File Formats)
#  Records (as large as possible)
#  Full Record and Cited References
#  Tab-delimited (Win)
# output:
#  JSON:
#   {
#     "author1": AU[0],
#     "year": PY,
#     "title": TI,
#     "abstract": AB,
#     "journal": J9,
#     "doi": DI,
#     "index": give a local number, 'newid' | 'tmpid'
#     "type": only choose article & review & editorial material
#     "cite": [ index1, index2, ..., indexN ]
#   }
# option:
#  all: all paper involved
#  rk1: besides the main paper list, paper cited by more than one main papers also included (# TODO)
#  min: only main papers
  '''
  fdata=codecs.open(filename,'rb',encoding);
  reader=csv.DictReader((uline.replace('\0','') for uline in fdata),delimiter='\t');
  odata=[]; # output
  tdata=[]; # temp stack
  newid=0; # output index
  tmpid=0; # temp stack index

  def __search_exists__(node):
    '''
# Search whether this node has been registered.
# return {'newid'|'tmpid':index} | None
    '''
    if (node.__contains__('doi') and node['doi']!=None and node['doi']!=''):
      index=None;
      for onode in odata:
        if (onode.__contains__('doi') and onode['doi']==node['doi']):
          index=onode['newid'];
          return {'newid':index};
      for tnode in tdata:
        if (tnode.__contains__('doi') and tnode['doi']==node['doi']):
          index=tnode['tmpid'];
          return {'tmpid':index};
    else:
      index=None;
      for onode in odata:
        if (onode.__contains__('doi') and onode['doi']!=None and onode['doi']!=''):
          continue;
        if (onode['author1']==node['author1'] and onode['year']==node['year'] and onode['journal']==node['journal']):
          index=onode['newid'];
          return {'newid':index};
      for tnode in tdata:
        if (tnode.__contains__('doi') and tnode['doi']!=None and tnode['doi']!=''):
          continue;
        if (tnode['author1']==node['author1'] and tnode['year']==node['year'] and tnode['journal']==node['journal']):
          index=tnode['tmpid'];
          return {'tmpid':index};
    return None;

  print('Loading Data ... Please Wait ...');
  for paper in reader: # load in data
    newnode={'author1':paper['AU'].split('; ')[0],'authors':paper['AU'],'title':paper['TI'],'abstract':paper['AB'],'journal':paper['J9'],'year':int(paper['PY']),'type':paper['DT'],'doi':paper['DI'],'citation':[]};
    thedoi=re.search('[0-9A-Za-z\./\-\%\(\)\[\]:]*',paper['DI']);
    if (thedoi):
      newnode['doi']=thedoi.group();
    for citestr in paper['CR'].split('; '):
      infos=citestr.split(', ');
      doi=None;
      for info in infos:
        thedoi=re.search('DOI [0-9A-Za-z\./\-\%\(\)\[\]:]*',info);
        if (thedoi):
          doi=thedoi.group()[4:];
          try:
           if (doi[0]=='['):
            doi=doi[1:];
          except Exception:
           print('Warning: Empty DOI. '+str(infos));
          break;
      if (len(infos)>2):
       try:
        tmpnode={'author1':infos[0],'year':int(infos[1]),'journal':infos[2],'doi':doi};
       except ValueError as err:
        print('Warning: citation contains invalid publish year data. where citation = '+str(infos));
        tmpnode={'author1':infos[0],'year':None,'journal':infos[1],'doi':doi};
      elif (len(infos)==2):
        tmpnode={'author1':infos[0],'doi':infos[1]};
        print('Warning: citation contains informal format. where citation = '+str(infos));
      elif (len(infos)==1):
        if (infos[0]!='' and infos[0]!=None):
          tmpnode={'doi':infos[0]};
          print('Warning: citation contains informal format. where citation = '+str(infos));
      idrst=__search_exists__(tmpnode);
      if (idrst==None):
        tmpnode['tmpid']=tmpid;
        tdata.append(tmpnode);
        newnode['citation'].append({'tmpid':tmpid});
        tmpid+=1;
      elif (idrst.__contains__('newid')):
        newnode['citation'].append({'newid':idrst['newid']}); # point at odata[idrst.newid]
      elif (idrst.__contains__('tmpid')):
        newnode['citation'].append({'tmpid':idrst['tmpid']}); # point at tdata[idrst.tmpid]
      else:
        print('Warning: __search_exists__ returns an invalid value.');
    newnode['newid']=newid;
    odata.append(newnode);
    newid+=1;
  for tndid in range(0,len(tdata)): # mark temp stack links
    if (tndid%100==0):
      print('Marking Links ... '+str(tndid)+'/'+str(len(tdata)));
    idrst=__search_exists__(tdata[tndid]);
    if (idrst==None):
      print('Warning: appended unfindable node ... where node = '+str(tdata[tndid]));
    elif (idrst.__contains__('newid')):
      tdata[tndid]['link']=idrst['newid'];
    else:
      if (option!='min'):
        tdata[tndid]['newid']=newid;
        tdata[tndid]['citation']=[];
        odata.append(tdata[tndid]);
        tdata[tndid]['link']=newid;
        newid+=1;
      else:
        tdata[tndid]['link']=None;
        # TODO: add rk1
  for ondid in range(0,len(odata)): # finally link the marks
    if (ondid%10==0):
      print('Linking Marks ... '+str(ondid)+'/'+str(len(odata)));
    cndid=0;
    while (cndid<len(odata[ondid]['citation'])):
      if (cndid%5==0):
        print('Running ... '+str(cndid)+'/'+str(len(odata[ondid]['citation'])));
      cnode=odata[ondid]['citation'][cndid];
      if (cnode.__contains__('tmpid')):
        if (option!='min'):
          newid=tdata[cnode['tmpid']]['link'];
          odata[ondid]['citation'][cndid]={'newid':newid};
        else:
          if (tdata[cnode['tmpid']]['link']!=None):
            newid=tdata[cnode['tmpid']]['link'];
            odata[ondid]['citation'][cndid]={'newid':newid};
            cndid+=1;
          else:
            odata[ondid]['citation'].pop(cndid);
            # TODO: add rk1
      else:
        cndid+=1;
  fdata.close();
  return odata;

